// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/simple_function_registry.h>

namespace kumo::pollux::exec {
    static SimpleFunctionRegistry instance;
    namespace {
        SimpleFunctionRegistry &simple_functions_internal() {
            return instance;
        }
    } // namespace

    const SimpleFunctionRegistry &simple_functions() {
        return simple_functions_internal();
    }

    SimpleFunctionRegistry &mutable_simple_functions() {
        return simple_functions_internal();
    }

    bool SimpleFunctionRegistry::register_function_internal(const std::string &name,
                                                          const std::shared_ptr<const Metadata> &metadata,
                                                          const FunctionFactory &factory, bool overwrite) {
        const auto sanitizedName = sanitize_name(name);
        return registeredFunctions_.withWLock([&](auto &map) {
            SignatureMap &signatureMap = map[sanitizedName];
            auto &functions = signatureMap[*metadata->signature()];

            for (auto it = functions.begin(); it != functions.end(); ++it) {
                const auto &otherMetadata = (*it)->getMetadata();

                if (metadata->physicalSignatureEquals(otherMetadata)) {
                    if (!overwrite) {
                        return false;
                    }
                    functions.erase(it);
                    break;
                }

                // Make sure default-null-behavior and deterministic properties are the
                // same for all implementations of a given logical signature.
                POLLUX_USER_CHECK_EQ(
                    metadata->defaultNullBehavior(), otherMetadata.defaultNullBehavior());
                POLLUX_USER_CHECK_EQ(
                    metadata->isDeterministic(), otherMetadata.isDeterministic());
            }
            functions.emplace_back(std::make_unique<const FunctionEntry>(metadata, factory));
            KLOG(INFO)<<"sanitizedName:"<<sanitizedName<<" functions: "<<functions.size()<<"map: "<<map.size();
            return true;
        });
    }

    void SimpleFunctionRegistry::removeFunction(const std::string &name) {
        const auto sanitizedName = sanitize_name(name);
        registeredFunctions_.withWLock([&](auto &map) { map.erase(sanitizedName); });
    }

    namespace {
        // This function is not thread safe. It should be called only from within a
        // synchronized read region of registeredFunctions_.
        const SignatureMap *getSignatureMap(
            const std::string &name,
            const FunctionMap &registeredFunctions) {
            const auto sanitizedName = sanitize_name(name);
            const auto it = registeredFunctions.find(sanitizedName);
            return it != registeredFunctions.end() ? &it->second : nullptr;
        }
    } // namespace

    std::vector<const FunctionSignature *>
    SimpleFunctionRegistry::get_function_signatures(const std::string &name) const {
        std::vector<const FunctionSignature *> signatures;
        registeredFunctions_.withRLock([&](const auto &map) {
            if (const auto *signatureMap = getSignatureMap(name, map)) {
                signatures.reserve(signatureMap->size());
                for (const auto &pair: *signatureMap) {
                    signatures.emplace_back(&pair.first);
                }
            }
        });

        return signatures;
    }

    std::vector<std::pair<VectorFunctionMetadata, const FunctionSignature *> >
    SimpleFunctionRegistry::get_function_signatures_and_metadata(
        const std::string &name) const {
        std::vector<std::pair<VectorFunctionMetadata, const FunctionSignature *> >
                result;
        registeredFunctions_.withRLock([&](const auto &map) {
            if (const auto *signatureMap = getSignatureMap(name, map)) {
                result.reserve(signatureMap->size());
                for (const auto &[signature, functions]: *signatureMap) {
                    VectorFunctionMetadata metadata{
                        false,
                        functions[0]->getMetadata().isDeterministic(),
                        functions[0]->getMetadata().defaultNullBehavior()
                    };
                    result.emplace_back(
                        std::pair<VectorFunctionMetadata, const FunctionSignature *>{
                            metadata, &signature
                        });
                }
            }
        });
        return result;
    }

    namespace {
        // Same as Type::kindEquals, but matches UNKNOWN in 'physicalType' to any type.
        bool physicalTypeMatches(const TypePtr &type, const TypePtr &physicalType) {
            if (physicalType->is_unKnown()) {
                return true;
            }

            if (type->kind() != physicalType->kind()) {
                return false;
            }

            if (type->size() != physicalType->size()) {
                return false;
            }

            for (auto i = 0; i < type->size(); ++i) {
                if (!physicalTypeMatches(type->childAt(i), physicalType->childAt(i))) {
                    return false;
                }
            }

            return true;
        }
    } // namespace

    std::optional<SimpleFunctionRegistry::ResolvedSimpleFunction>
    SimpleFunctionRegistry::resolve_function(
        const std::string &name,
        const std::vector<TypePtr> &argTypes) const {
        const FunctionEntry *selectedCandidate = nullptr;
        TypePtr selectedCandidateType = nullptr;
        registeredFunctions_.withRLock([&](const auto &map) {
            if (const auto *signatureMap = getSignatureMap(name, map)) {
                for (const auto &[candidateSignature, functionEntry]: *signatureMap) {
                    SignatureBinder binder(candidateSignature, argTypes);
                    if (binder.tryBind()) {
                        for (const auto &currentCandidate: functionEntry) {
                            const auto &m = currentCandidate->getMetadata();

                            // For variadic signatures, number of arguments in function call may
                            // be one less than number of arguments in the signature.
                            const auto numArgsToMatch =
                                    std::min(argTypes.size(), m.argPhysicalTypes().size());

                            bool match = true;
                            for (auto i = 0; i < numArgsToMatch; ++i) {
                                if (!physicalTypeMatches(argTypes[i], m.argPhysicalTypes()[i])) {
                                    match = false;
                                    break;
                                }
                            }

                            if (!match) {
                                continue;
                            }

                            if (!selectedCandidate ||
                                currentCandidate->getMetadata().priority() <
                                selectedCandidate->getMetadata().priority()) {
                                auto resultType = binder.tryResolveReturnType();
                                POLLUX_CHECK_NOT_NULL(resultType);

                                if (physicalTypeMatches(resultType, m.resultPhysicalType())) {
                                    selectedCandidate = currentCandidate.get();
                                    selectedCandidateType = resultType;
                                }
                            }
                        }
                    }
                }
            }
        });

        POLLUX_DCHECK(!selectedCandidate || selectedCandidateType);

        return selectedCandidate
                   ? std::optional<ResolvedSimpleFunction>(
                       ResolvedSimpleFunction(*selectedCandidate, selectedCandidateType))
                   : std::nullopt;
    }
} // namespace kumo::pollux::exec
